# -*- coding: utf-8 -*-

"""最小生成树算法之 prim，详解见书"""

class Prim():
    def __init__(self) -> None:
        pass

    def mini_spanning_tree(self, vertex, edge, w):
        """input:
        vertex:= {a,b,c,d,e}
        edge:= {{a,b}, {d,c}, {b,e}, {c,e}}
        w:= f: edge -> number
        -------------------------
        mini_tree_vertex:= {a, b} 表示已经产生的最小生成树子树的顶点
        mini_tree_edge:= {{a,b}} 表示已经产生的最小生成树子树的边"""

        mini_tree_vertex, mini_tree_edge = set(), set()
        mini_tree_vertex.add(vertex.pop())

        # 可直接用原始的顶点集表示生成子树顶点的补集，因为顶点不断在从顶点补集中迁移到顶点生成集
        # 可直接用原始的边集表示生成树边的补集，因为边也不断在边补集中迁移到生成边集

        while vertex:
            u1, v1, e12, w12 = None, None, None, float('inf')
            for u in mini_tree_vertex:    # 子树的顶点
                for v in vertex:    # 子树以外的顶点
                    e = frozenset((u, v))
                    if e in edge and w(e) < w12:    # 内外顶点形成的边在子树边的补集中 且 该边的权要小一些
                        u1, v1, e12, w12 = u, v, e, w(e)
            
            # 找到子树顶点补集中的v，以及连接内外的边，更新相应容器
            vertex.discard(v1)
            mini_tree_vertex.add(v1)
            edge.discard(e12)
            mini_tree_edge.add(e12)

        return mini_tree_edge


if __name__ == '__main__':
    vertex = {'a','b','c','d','e','f'}
    edge = {frozenset(('a','b'))
            , frozenset(('d','e'))
            , frozenset(('e','f'))
            , frozenset(('a','d'))
            , frozenset(('b','e'))
            , frozenset(('c','f'))
            , frozenset(('a','e'))
            , frozenset(('b','d'))
            , frozenset(('c','e'))}
    def w(e):
        weight = {
            frozenset(('a','b')): 1
            , frozenset(('d','e')): 4
            , frozenset(('e','f')): 7
            , frozenset(('a','d')): 4
            , frozenset(('b','e')): 2
            , frozenset(('c','f')): 5
            , frozenset(('a','e')): 3
            , frozenset(('b','d')): 3
            , frozenset(('c','e')): 4
        }
        return weight[e]
    
    prim = Prim()
    mini_tree = prim.mini_spanning_tree(vertex, edge, w)
    print(mini_tree)

